package org.example.elasticsearch.query;


import co.elastic.clients.elasticsearch._types.*;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.query_dsl.BoolQuery;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.search.SourceConfig;
import co.elastic.clients.elasticsearch.core.search.SourceFilter;
import co.elastic.clients.json.JsonData;
import co.elastic.clients.util.ObjectBuilder;
import org.example.domain.PageQuery;
import org.example.elasticsearch.ESIndex;
import org.example.elasticsearch.ESIndexEnum;
import org.example.elasticsearch.ESUtil;
import org.example.utils.Func;
import lombok.Getter;

import java.io.Serializable;
import java.util.*;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * @author: sql
 * @date: 2024/6/7
 * @Description: ES查询参数Builder类
 * @version: 1.0
 */
public class ESQueryBuilder {

    private final List<String> indexNames = new ArrayList<>();

    @Getter
    private ESIndex esIndex;

    private Class<?> sourceClass;

    private int pageNum = 1;
    private int pageSize = -1;

    private List<SortOptions> sorts;

    private List<String> includesShowFields;
    private List<String> excludesShowFields;

    private Map<String, Aggregation> aggregations;

    private final BoolQuery.Builder queryBuilder;

    private ESQueryWrapper must;
    private ESQueryShouldWrapper should;
    private ESQueryWrapper filter;
    private ESQueryWrapper mustNot;
    private Integer preciseCount;

    private Map<String, ScriptField> scriptFieldMap;

    private ESQueryBuilder(String indexName, String... indexNames) {
        this.indexNames.add(indexName);
        if (null != indexNames) {
            this.indexNames.addAll(Arrays.asList(indexNames));
        }
        this.queryBuilder = new BoolQuery.Builder();
    }

    public ESQueryBuilder(ESIndex esIndex) {
        this.esIndex = esIndex;
        indexNames.add(ESUtil.getIndexName(esIndex));
        this.queryBuilder = new BoolQuery.Builder();
    }


    /**
     * 获取ESQueryBuilder，传入要搜索的索引名
     *
     * @param indexName  必传的索引
     * @param indexNames 选传的其他索引，一起搜索？
     * @return ESQueryBuilder，记得最后要build()，获得SearchRequest对象
     */
    public static ESQueryBuilder of(String indexName, String... indexNames) {
        return new ESQueryBuilder(indexName, indexNames);
    }

    /**
     * 获取ESQueryBuilder，传入要搜索的索引名
     *
     * @param indexEnum 必传的索引枚举
     * @return ESQueryBuilder，记得最后要build()，获得SearchRequest对象
     */
    public static ESQueryBuilder of(ESIndexEnum indexEnum) {
        return new ESQueryBuilder(indexEnum);
    }


    /**
     * 整理所有参数，最终调用这个方法获取SearchRequest对象
     *
     * @return 构建出SearchRequest对象
     */
    public SearchRequest build() {

        if (!(null == must || must.getQueries().isEmpty())) {
            queryBuilder.must(must.getQueries());
        }

        if (!(null == should || should.getQueries().isEmpty())) {
            queryBuilder.should(should.getQueries());
            if (Func.isNotEmpty(should.getMinShouldMatch())) {
                queryBuilder.minimumShouldMatch(should.getMinShouldMatch());
            }
        }

        if (!(null == filter || filter.getQueries().isEmpty())) {
            queryBuilder.filter(filter.getQueries());
        }

        if (!(null == mustNot || mustNot.getQueries().isEmpty())) {
            queryBuilder.mustNot(mustNot.getQueries());
        }

        BoolQuery query = queryBuilder.build();

        Function<SearchRequest.Builder, ObjectBuilder<SearchRequest>> builder = getBuilderFunction(query);

        return SearchRequest.of(builder);
    }

    /**
     * bool查询中的Must场景
     * 设置的所有参数必须全部满足才会匹配到
     *
     * @param func must场景下的各种查询参数设置
     * @return 自己
     */
    public ESQueryBuilder must(Consumer<ESQueryWrapper> func) {
        if (null == must) {
            must = new ESQueryWrapper();
        }
        func.accept(must);
        return this;
    }


    /**
     * bool查询中的should场景
     * 设置的所有参数只要有满足的就会匹配到
     *
     * @param func should场景下的各种查询参数设置
     * @return 自己
     */
    public ESQueryBuilder should(Consumer<ESQueryShouldWrapper> func) {
        if (null == should) {
            should = new ESQueryShouldWrapper();
        }
        func.accept(should);
        return this;
    }


    /**
     * bool查询中的filter场景
     * 设置的所有参数只要有满足的就会匹配到
     *
     * @param func filter场景下的各种查询参数设置
     * @return 自己
     */
    public ESQueryBuilder filter(Consumer<ESQueryWrapper> func) {
        if (null == filter) {
            filter = new ESQueryWrapper();
        }
        func.accept(filter);
        return this;
    }

    /**
     * bool查询中的mustNot场景
     * 设置的所有参数只要有满足的就会匹配到
     *
     * @param func mustNot场景下的各种查询参数设置
     * @return 自己
     */
    public ESQueryBuilder mustNot(Consumer<ESQueryWrapper> func) {
        if (null == mustNot) {
            mustNot = new ESQueryWrapper();
        }
        func.accept(mustNot);
        return this;
    }

    /**
     * 分页参数
     *
     * @param pageNum  当前页数，从1开始
     * @param pageSize 每页多少条
     * @return 自己
     */
    public ESQueryBuilder setPage(int pageNum, int pageSize) {
        this.pageNum = Math.max(pageNum, 1);
        this.pageSize = Math.max(pageSize, 1);
        return this;
    }

    public <T extends PageQuery> ESQueryBuilder setPage(T pageQuery) {
        this.pageNum = (int) Math.max(pageQuery.getPageNumber(), 1);
        this.pageSize = (int) Math.max(pageQuery.getPageSize(), 1);
        return this;
    }

    /**
     * 可以自定义查询字段
     *
     * @param sourceClass
     * @return
     */
    public ESQueryBuilder sourceClass(Class<?> sourceClass) {
        this.sourceClass = sourceClass;
        return this;
    }


    /**
     * 排序
     *
     * @param field 要排序的字段
     * @param desc  是否降序
     * @return 自己
     */
    public ESQueryBuilder sort(String field, boolean desc) {
        if (null == sorts) {
            sorts = new LinkedList<>();
        }
        sorts.add(
                SortOptions.of(s ->
                        s.field(
                                f -> f.field(field)
                                        .order(desc ? SortOrder.Desc : SortOrder.Asc)
                        )
                )
        );

        return this;
    }

    /**
     * 排序，将特殊值修改为指定值
     *
     * @param field 要排序的字段
     * @param desc  是否降序
     * @return 自己
     */
    public ESQueryBuilder sortWithSpecialValues(String field, boolean desc, Map<String, Integer> specialValueMap) {
        if (null == sorts) {
            sorts = new LinkedList<>();
        }

        Map<String, JsonData> param = Func.newHashMap();
        param.put("field", JsonData.of(field));
        specialValueMap.forEach((k, v) -> param.put(k, JsonData.of(v)));


        sorts.add(
                SortOptions.of(b -> b.script(s -> s.script(sb -> sb.inline(in ->
                                in.source(
                                                """
                                                            
                                                            String key = String.valueOf(doc[params.field].value);
                                                            
                                                            if (params.containsKey(key)) {
                                                                return params[key];
                                                            }
                                                            
                                                            return Integer.parseInt(doc[params.field].value);
                                                            
                                                        """)
                                        .params(param)
                                        .lang("painless"))).type(ScriptSortType.Number)
                        .order(desc ? SortOrder.Desc : SortOrder.Asc)))
        );

        return this;
    }


    public ESQueryBuilder sortWithSpecialValues(String field, boolean desc, String... specialValues) {
        Map<String, Integer> specialValueMap = new HashMap<>();
        int index = desc ? Integer.MAX_VALUE : Integer.MIN_VALUE;
        for (String v : specialValues) {
            specialValueMap.put(v, desc ? index-- : index++);
        }
        return sortWithSpecialValues(field, desc, specialValueMap);
    }


    /**
     * 全部升序排序
     *
     * @param field 要升序排序的字段
     * @return 自己
     */
    public ESQueryBuilder sortAsc(String... field) {
        if (null == sorts) {
            sorts = new LinkedList<>();
        }

        for (String item : field) {
            sorts.add(
                    SortOptions.of(s ->
                            s.field(
                                    f -> f.field(item)
                                            .order(SortOrder.Asc)
                            )
                    )
            );
        }


        return this;
    }


    /**
     * 全部降序排序
     *
     * @param field 要降序排序的字段
     * @return 自己
     */
    public ESQueryBuilder sortDesc(String... field) {
        if (null == sorts) {
            sorts = new LinkedList<>();
        }

        for (String item : field) {
            sorts.add(
                    SortOptions.of(s ->
                            s.field(
                                    f -> f.field(item)
                                            .order(SortOrder.Desc)
                            )
                    )
            );
        }

        return this;
    }


    /**
     * 最外层的聚合控制
     * 对所有匹配条件命中的数据进行聚合统计
     *
     * @param func 聚合的具体操作
     * @return 自己
     */
    public ESQueryBuilder aggregations(Consumer<ESAggregationWrapper> func) {
        if (null == aggregations) {
            aggregations = new HashMap<>();
        }

        func.accept(ESAggregationWrapper.get(aggregations));

        return this;
    }


    /**
     * 分组，并对组内进行聚合操作
     *
     * @param showName     组名
     * @param computeField 要分组的字段
     * @param func         具体组内聚合操作
     * @return 自己
     */
    public ESQueryBuilder groupAggregations(String showName, String computeField, Consumer<ESAggregationWrapper> func) {
        if (null == aggregations) {
            aggregations = new HashMap<>();
        }

        ESAggregationWrapper.get(aggregations)
                .groupByFields(showName, computeField, func);

        return this;
    }


    /**
     * 返回值中需要包含的字段
     *
     * @param fields 需要包含的字段们
     * @return 自己
     */
    public ESQueryBuilder addIncludesShowFields(String... fields) {
        if (null == includesShowFields) {
            includesShowFields = new ArrayList<>();
        }
        includesShowFields.addAll(Arrays.asList(fields));
        return this;
    }


    /**
     * 返回值中需要排除的字段
     *
     * @param fields 需要排除的字段们
     * @return 自己
     */
    public ESQueryBuilder addExcludesShowFields(String... fields) {
        if (null == excludesShowFields) {
            excludesShowFields = new ArrayList<>();
        }
        excludesShowFields.addAll(Arrays.asList(fields));
        return this;
    }


    /**
     * 开启精准计数，返回精准count，上限为100W
     *
     * @return 自己
     */
    public ESQueryBuilder preciseCount() {
        this.preciseCount = 1_000_000;
        return this;
    }


    /**
     * 开启精准计数，返回精准count，上限自己控制
     *
     * @return 自己
     */
    public ESQueryBuilder preciseCount(int max) {
        this.preciseCount = max;
        return this;
    }


    // ================================= 下面的不用管 =================================

    private Function<SearchRequest.Builder, ObjectBuilder<SearchRequest>> getBuilderFunction(BoolQuery query) {

        Function<SearchRequest.Builder, ObjectBuilder<SearchRequest>> builder =
                b -> b.index(indexNames).timeout("10s").query(q -> q.bool(query));

        if (pageSize > 0) {
            builder = builder
                    .andThen(b -> ((SearchRequest.Builder) b)
                            .from((pageNum - 1) * pageSize)
                            .size(pageSize));
        }

        if (!(sorts == null || sorts.isEmpty())) {
            builder = builder
                    .andThen(b -> ((SearchRequest.Builder) b)
                            .sort(sorts));
        }

        if (!(aggregations == null || aggregations.isEmpty())) {
            builder = builder
                    .andThen(b -> ((SearchRequest.Builder) b).aggregations(aggregations));
        }

        if (null != preciseCount) {
            builder = builder
                    .andThen(b -> ((SearchRequest.Builder) b).trackTotalHits(h -> h.count(100000)));
        }


        if (null != includesShowFields || null != excludesShowFields) {
            builder = builder
                    .andThen(b -> ((SearchRequest.Builder) b).source(s -> s.filter(f -> {
                        SourceFilter.Builder res = null;
                        if (null != includesShowFields) {
                            res = f.includes(includesShowFields);
                        }
                        if (null != excludesShowFields) {
                            res = f.excludes(excludesShowFields);
                        }
                        return res;
                    })));
        }

        if (null != scriptFieldMap) {
            builder = builder.andThen(b -> {
                scriptFieldMap.forEach((k, v) -> ((SearchRequest.Builder) b).scriptFields(k, v));
                return b;
            });
        }

        if (null != esIndex) {
            builder = builder.andThen(b -> {
                //可以自定义查询类
                Class<?> esDocClass = esIndex.getEsDocClass();
                if (sourceClass != null) {
                    esDocClass = sourceClass;
                }
                List<String> fieldList = Func.getFieldListNames(esDocClass);
                ((SearchRequest.Builder) b).source(SourceConfig.of(s -> s.filter(v -> v.includes(fieldList))));
                return b;
            });
        }

        return builder;
    }

    /**
     * TODO 王老板 来写
     *
     * @param fieldName
     * @param scriptField
     * @return
     */
    public ESQueryBuilder scriptFields(String fieldName, ScriptField scriptField) {

        if (null == scriptFieldMap) scriptFieldMap = new HashMap<>();
        scriptFieldMap.put(fieldName, scriptField);

        return this;
    }


    @Getter
    public static class ESQueryWrapper {

        private final LinkedList<Query> queries = new LinkedList<>();

        // TODO 测试下，多个term会不会自动变成terms
        public ESQueryWrapper term(boolean expression, String field, Serializable value) {
            if (expression) {
                term(field, value);
            }
            return this;
        }

        public ESQueryWrapper term(String field, Serializable value) {
            if (Func.isNotEmpty(value)) {
                queries.add(Query.of(q -> q.term(x -> x.field(field).value(value.toString()))));
            }
            return this;
        }

        /**
         * 满足集合中的数据就查询出来
         *
         * @param field
         * @param values
         * @return
         */
        public ESQueryWrapper terms(String field, List<?> values) {
            if (Func.isNotEmpty(values)) {
                List<FieldValue> list = values.stream().map(FieldValue::of).toList();
                queries.add(Query.of(q -> q.terms(x -> x.field(field).terms(t -> t.value(list)))));
            }
            return this;
        }

        public ESQueryWrapper match(String field, Serializable value) {
            if (Func.isNotEmpty(value)) {
                queries.add(Query.of(q -> q.match(x -> x.field(field).query(value.toString()))));
            }
            return this;
        }

        public ESQueryWrapper match(Serializable value, String... fields) {
            if (Func.isNotEmpty(value) && Func.isNotEmpty(fields)) {
                for (String field : fields) {
                    queries.add(Query.of(q -> q.match(x -> x.field(field).query(value.toString()))));
                }
            }
            return this;
        }

        public ESQueryWrapper regexp(Serializable value, String... fields) {
            if (Func.isNotEmpty(value) && Func.isNotEmpty(fields)) {
                for (String field : fields) {
                    queries.add(Query.of(q -> q.regexp(x -> x.field(field).value(value.toString()))));
                }
            }
            return this;
        }

        public ESQueryWrapper wildcard(Serializable value, String... fields) {
            if (Func.isNotEmpty(value) && Func.isNotEmpty(fields)) {
                for (String field : fields) {
                    queries.add(Query.of(q -> q.wildcard(x -> x.field(field).value(String.format("*%s*", value)))));
                }
            }
            return this;
        }

        public ESQueryWrapper like(Serializable value, String... fields) {
            return wildcard(value, fields);
        }

        public ESQueryWrapper gt(boolean expression, String field, Serializable val) {
            if (expression) {
                gt(field, val);
            }
            return this;
        }

        public ESQueryWrapper gt(String field, Serializable val) {
            if (Func.isNotEmpty(val)) {
                queries.add(Query.of(q -> q.range(x -> {
                    x.field(field);
                    x.gt(JsonData.of(val));
                    return x;
                })));
            }
            return this;
        }

        /**
         * 范围查询
         *
         * @param field
         * @param from
         * @param to
         * @return
         */
        public ESQueryWrapper range(String field, Serializable from, Serializable to) {
            if (Func.isNotEmpty(from) || Func.isNotEmpty(to)) {
                queries.add(Query.of(q -> q.range(x -> {
                    x.field(field);
                    if (Func.isNotEmpty(from)) {
                        x.from(from.toString());
                    }
                    if (Func.isNotEmpty(to)) {
                        x.to(to.toString());
                    }
                    return x;
                })));
            }
            return this;
        }

        /**
         * 查询天数
         *
         * @param field
         * @param state 查询天数（1：近一天，2：近一周，3：近一个月，4：近三个月）
         * @return
         */
        public ESQueryWrapper rangeDay(String field, String state) {
            if (Func.isNotEmpty(state)) {
                String gte = switch (state) {
                    case "1" -> "now-1d/d";//近一天
                    case "2" -> "now-1w/d";//近一周
                    case "3" -> "now-1M/d";//近一个月
                    case "4" -> "now-3M/d";//近三个月
                    default -> null;
                };
                if (Func.isNotEmpty(gte)) {
                    queries.add(Query.of(q -> q.range(x -> x.field(field).gte(JsonData.of(gte)).lt(JsonData.of("now+1d/d")))));
                }
            }
            return this;
        }

        public ESQueryWrapper fuzzy(String field, Serializable value) {
            if (Func.isNotEmpty(value)) {
                queries.add(Query.of(q -> q.fuzzy(x -> x.field(field).value(value.toString()))));
            }
            return this;
        }

    }

    @Getter
    public static class ESAggregationWrapper {

        private final Map<String, Aggregation> myAggregations;

        public static ESAggregationWrapper get() {
            return new ESAggregationWrapper(new HashMap<>());
        }

        public static ESAggregationWrapper get(Map<String, Aggregation> myAggregations) {
            return new ESAggregationWrapper(myAggregations);
        }


        private ESAggregationWrapper(Map<String, Aggregation> aggregations) {
            this.myAggregations = aggregations;
        }

        public ESAggregationWrapper max(String computeField) {
            myAggregations.put("max_" + computeField, Aggregation.of(a -> a.max(m -> m.field(computeField))));
            return this;
        }

        public ESAggregationWrapper max(String showName, String computeField) {
            myAggregations.put(showName, Aggregation.of(a -> a.max(m -> m.field(computeField))));
            return this;
        }

        public ESAggregationWrapper min(String computeField) {
            myAggregations.put("min_" + computeField, Aggregation.of(a -> a.min(m -> m.field(computeField))));
            return this;
        }

        public ESAggregationWrapper min(String showName, String computeField) {
            myAggregations.put(showName, Aggregation.of(a -> a.min(m -> m.field(computeField))));
            return this;
        }

        public ESAggregationWrapper sum(String computeField) {
            myAggregations.put("sum_" + computeField, Aggregation.of(a -> a.sum(m -> m.field(computeField))));
            return this;
        }

        public ESAggregationWrapper sum(String showName, String computeField) {
            myAggregations.put(showName, Aggregation.of(a -> a.sum(m -> m.field(computeField))));
            return this;
        }

        public ESAggregationWrapper groupByFields(String computeField) {
            return groupByFields("group_" + computeField, computeField, null);
        }

        public ESAggregationWrapper groupByFields(String showName, String computeField) {
            return groupByFields(showName, computeField, null);
        }

        public ESAggregationWrapper groupByFields(String showName, String computeField,
                                                  Consumer<ESAggregationWrapper> func) {

            if (null == func) {
                myAggregations.put(showName, Aggregation.of(
                        a -> a.terms(t -> t.field(computeField)))
                );
            } else {

                ESAggregationWrapper esAggregationWrapper = ESAggregationWrapper.get();

                func.accept(esAggregationWrapper);

                myAggregations.put(showName, Aggregation.of(
                        a -> a.terms(t -> t.field(computeField))
                                .aggregations(esAggregationWrapper.getMyAggregations()))
                );

            }

            return this;
        }


    }


}
